What did I do this week?
Following last week’s investigation into Issue #47, where the JOPT algorithm fails for open-system state transfer using the TRACEDIFF
fidelity type, I revisited and reworked the proposed solution based on mentor feedback.
In my initial PR, I attempted to fix the bug by manually constructing the adjoint using:
=Dimensions(diff.dims[1], diff.dims[0]), copy=False) Qobj(diff.data.adjoint(), dims
During our weekly meeting, we discussed that the Dimensions
class doesn’t exist in the version of QuTiP (v5.1) used in qutip-qoc
. This prompted us to explore simpler, version-compatible alternatives.
🧠 Mentor Suggestions
@BoxiLi suggested that since diff
is already a Qobj
, we can simply take the adjoint using:
= diff.dag() diff_dag
Alex agreed with this direction, noting that “perhaps all we need is an appropriately placed .dag()
”.
🔍 What I Tried
We replaced the manual adjoint with:
= diff.dag()
diff_dag = 0.5 * (diff_dag * diff).tr() g
However, this caused compatibility issues with JAX. Specifically, calling .tr()
raised an AttributeError
because the underlying object became a DynamicJaxprTracer
, which doesn’t support typical QuTiP methods like .tr()
, .full()
, or .data.trace()
.
🛠️ What I am Doing Now
To address the issue, I am now exploring a JAX-compatible alternative that uses raw JAX array operations instead of QuTiP methods. For example:
import jax.numpy as jnp
= 0.5 * jnp.trace(jnp.matmul(diff_dag._jxa, diff._jxa)) g
Our goal is to: - ✅ Maintain a minimal patch size (as mentors prefer) - 🔄 Avoid QuTiP internals that break with JAX tracing - ⚙️ Preserve logic correctness for density matrices in Liouville space
🧪 Testing
The updated fix is being tested on a minimal open-system JOPT use case that previously failed with TRACEDIFF. The test is verifying that:
diff.dag()
behaves as expected- The fidelity value is computed via JAX-compatible expressions
- The new logic doesn’t introduce new shape or gradient issues
🔧 Experimental Qtrl Integration
I also began early experiments to explore the dependency and fidelity logic in qutip-qtrl
. This included:
- Identifying key parts of
qtrl
that overlap withqutip-qoc
- Experimentally copying selected fidelity-related files from
qtrl
into thequtip-qoc
codebase - Testing whether they can be made to work as a proof-of-concept, potentially as part of a future unification
This is just a trial, but it helps identify where tight coupling exists and what effort would be required to integrate or migrate logic more formally.
Plan for next week
I’m excited to be moving closer to a JAX-compatible, open-system-aware fidelity framework that can serve as a reliable core for QuTiP-QOC’s optimization methods.